package developer.zyk.ability.chat.template;

import developer.zyk.ability.chat.dao.entity.ChatHistoryDO;
import developer.zyk.ability.chat.param.ChatAppendMessageParam;
import developer.zyk.ability.chat.service.ChatHistoryService;
import developer.zyk.ability.helper.ModelConfigHelper;
import developer.zyk.global.user.bo.config.ModelConfig;
import developer.zyk.global.util.AssertUtil;
import jakarta.annotation.Resource;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.QuestionAnswerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Flux;

import java.util.Objects;

/**
 * @author zhangyongkang
 * @time 2025/5/15 20:05
 * @description
 */
@Component
public class ChatTemplate {

    @Resource
    private ChatClient chatClient;
    @Resource
    private ChatHistoryService chatHistoryService;
    @Resource
    private ChatMemory dbChatMemory;
    @Resource
    private VectorStore vectorStore;


    public Flux<String> appendMessage(ChatAppendMessageParam param) {
        //step1 获取基础信息
        Long chatId = param.getChatId();
        String userMessage = param.getMessage();
        ModelConfig modelConfig = ModelConfigHelper.getModelConfigByUserId(param.getUserId());
        ChatHistoryDO chatInfo = getChatHistory(param.getChatId());
        String modelType = chatInfo.getModelType();//当前使用模型

        ModelConfig.ModelConfigDetail modelConfigDetail = modelConfig.getModelConfig(modelType);
        AssertUtil.assertNonNull(modelConfigDetail, "会话配置为空");
        //step2 构建自定义的聊天客户端
        ModelConfig.ModelConfigDetail.RagConfig ragConfig = modelConfigDetail.getRagConfig();

        //2.1 系统提示词
        ChatClient.ChatClientRequestSpec chatClientBuilder = chatClient
                .prompt()
                .system(ModelConfigHelper.systemPromptTemplate(modelConfigDetail::getSystemPromptTemplate));
        //2.2 上下文长度
        chatClientBuilder.advisors(MessageChatMemoryAdvisor.builder(dbChatMemory)
                .conversationId(chatId.toString())
                .chatMemoryRetrieveSize(
                        ModelConfigHelper.historySize(modelConfigDetail::getHistorySize)
                )
                .build());
        //2.3 rag配置
        if (Objects.equals(ModelConfigHelper.ragEnable(ragConfig::getEnableRag), Boolean.TRUE)) {
            chatClientBuilder.advisors(
                    new QuestionAnswerAdvisor(vectorStore,
                            SearchRequest.builder()
                                    .topK(ModelConfigHelper.ragTopN(ragConfig::getTopN))
                                    .similarityThreshold(ModelConfigHelper.similarityThreshold(ragConfig::getSimilarityThreshold))
                                    .build(),
                            ModelConfigHelper.ragTemplate(ragConfig::getPromptTemplate))
            );
        }
        //2.3 用户信息
        chatClientBuilder.user(userMessage);

        //step 3 流式返回
        return chatClientBuilder
                .options(ChatOptions.builder()
                        .model(modelType)
                        .frequencyPenalty(
                                ModelConfigHelper.frequencyPenalty(modelConfigDetail::getFrequencyPenalty)
                        )
                        .build())
                .stream()
                .content();
    }

    private ChatHistoryDO getChatHistory(Long chatId) {
        return chatHistoryService.findChatByChatId(chatId);
    }
}
